import numpy as np
import cv2
import yaml
from collections import deque


def floating_center_discretize(img, runnable_color=100, node_radius=30):
    h, w = img.shape

    # 1. 收集所有可行像素并打乱
    print("[1/5] 正在收集可行像素...")
    r_coords = np.argwhere(img == runnable_color)
    r_list = [tuple(pt) for pt in r_coords]
    np.random.shuffle(r_list)

    # 2. 准备两个掩码
    print("[2/5] 正在准备掩码...")
    temp_visited = np.zeros((h, w), dtype=np.bool_)
    visited = np.zeros((h, w), dtype=np.bool_)

    nodes = []  # 最终浮动中心列表（存 [row_mean, col_mean]）
    area_list = []  # 每个区域的像素集合

    # 3. 依次遍历打乱后的可行像素，若未被 visited，就当作新的“种子”启动区域生长
    print("[3/5] 正在遍历像素...")
    for r0, c0 in r_list:
        if visited[r0, c0]:
            continue

        print("[3/5] 正在处理像素：(%d, %d)" % (r0, c0))
        queue = deque([(r0, c0)])
        temp_visited[r0, c0] = True
        point_list = list(queue)  # 初始化点列表

        # 初始浮动中心
        center = np.array([r0, c0], dtype=np.float64)
        new_points_count = 0

        while queue:
            pr, pc = queue.popleft()

            for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                nr, nc = pr + dr, pc + dc
                if not (0 <= nr < h and 0 <= nc < w):
                    continue
                if temp_visited[nr, nc]:
                    continue
                if img[nr, nc] == runnable_color:
                    dx = abs(nr - center[0])
                    dy = abs(nc - center[1])
                    if dx <= node_radius and dy <= node_radius:
                        temp_visited[nr, nc] = True
                        queue.append((nr, nc))
                        point_list.append((nr, nc))
                        new_points_count += 1

            # 每新增10个点更新一次浮动中心
            if new_points_count >= 10:
                center = np.mean(point_list, axis=0)
                new_points_count = 0

        # 精修区域，剔除离最终中心过远的像素
        final_center = np.mean(point_list, axis=0)
        valid_pixels = []
        for pr, pc in point_list:
            dx = abs(pr - final_center[0])
            dy = abs(pc - final_center[1])
            if dx <= node_radius and dy <= node_radius:
                valid_pixels.append((pr, pc))
            else:
                temp_visited[pr, pc] = False

        for pr, pc in valid_pixels:
            visited[pr, pc] = True

        nodes.append(final_center.tolist())
        area_list.append(valid_pixels)

    # 新增步骤：合并小面积区域
    print("[新增步骤] 正在合并小面积区域...")

    # 1. 计算面积和均值
    areas_size = [len(area) for area in area_list]
    total_size = sum(areas_size)
    mean_size = total_size / len(areas_size) if len(areas_size) > 0 else 0
    threshold = mean_size * 0.2
    print("  >> 平均区域面积: %.2f 像素" % mean_size)
    print("  >> 小面积阈值: %.2f 像素" % threshold)

    # 2. 筛选小区域并按面积排序
    small_areas = [i for i in range(len(areas_size)) if areas_size[i] < threshold]
    small_areas.sort(key=lambda x: areas_size[x])
    print("  >> 待合并区域数量: %d" % len(small_areas))

    # 3. 构建临时标签图用于邻接检测
    label_temp = np.full((h, w), -1, dtype=np.int32)
    for idx, pixels in enumerate(area_list):
        for r, c in pixels:
            label_temp[r, c] = idx

    # 4. 预计算邻接表
    print("  >> 正在预计算邻接表...")
    adjacency = [[] for _ in range(len(area_list))]
    offsets = np.array([[-1, 0], [1, 0], [0, -1], [0, 1]])
    for idx, pixels in enumerate(area_list):
        coords = np.array(pixels)
        neighbor_coords = coords[:, None, :] + offsets[None, :, :]
        valid_mask = (
            (neighbor_coords[:, :, 0] >= 0)
            & (neighbor_coords[:, :, 0] < h)
            & (neighbor_coords[:, :, 1] >= 0)
            & (neighbor_coords[:, :, 1] < w)
        )
        neighbor_coords = neighbor_coords[valid_mask]
        labels = label_temp[neighbor_coords[:, 0], neighbor_coords[:, 1]]
        unique_labels = np.unique(labels[labels != -1])
        adjacency[idx] = [int(l) for l in unique_labels if l != idx]

    # 5. 处理每个小区域，合并到最近邻接区域
    merged = [False] * len(area_list)  # 标记是否已被合并
    centers = np.array(nodes)

    print("  >> 正在合并区域...")
    for i in small_areas:
        if merged[i]:
            continue

        neighbors = adjacency[i]
        valid_candidates = [j for j in neighbors if not merged[j] and j != i]

        if not valid_candidates:
            print("  >> 区域 %d 无邻接区域，跳过合并。" % i)
            continue

        # 向量计算距离
        dists = np.linalg.norm(centers[i] - centers[valid_candidates], axis=1)
        j_min = valid_candidates[np.argmin(dists)]

        print(
            "  >> 区域 %d 合并入区域 %d，距离: %.2f"
            % (i, j_min, dists[np.argmin(dists)])
        )

        # 合并区域 i 到 j_min
        area_list[j_min].extend(area_list[i])
        centers[j_min] = np.mean(area_list[j_min], axis=0)
        merged[i] = True

    # 6. 过滤掉已被合并的区域，更新 area_list 和 nodes
    new_area_list = []
    new_nodes = []

    for idx in range(len(area_list)):
        if not merged[idx]:
            new_area_list.append(area_list[idx])
            new_nodes.append(centers[idx].tolist())

    area_list = new_area_list
    nodes = new_nodes
    print("  >> 合并完成，剩余区域数量: %d" % len(area_list))

    # 4. 构造 label 图，并一次性扫描找出所有连边
    print("[4/5] 正在构造 label 图...")
    label_img = -1 * np.ones((h, w), dtype=np.int32)
    for idx, pixs in enumerate(area_list):
        for pr, pc in pixs:
            label_img[pr, pc] = idx

    edge_set = set()
    for r in range(h):
        for c in range(w):
            li = label_img[r, c]
            if li < 0:
                continue
            if c + 1 < w:
                lj = label_img[r, c + 1]
                if lj >= 0 and lj != li:
                    edge_set.add((min(li, lj), max(li, lj)))
            if r + 1 < h:
                lj = label_img[r + 1, c]
                if lj >= 0 and lj != li:
                    edge_set.add((min(li, lj), max(li, lj)))

    edges = [list(e) for e in edge_set]
    edges = [[int(e[0]), int(e[1])] for e in edges]

    # 5. 绘图：灰度图转 BGR + 区域染色 + 圆点 + 连边
    print("[5/5] 正在进行可视化处理")
    img_color = cv2.cvtColor(img.copy(), cv2.COLOR_GRAY2BGR)

    # 每个区域染色
    for pixs in area_list:
        color = np.random.randint(0, 255, size=3).tolist()
        rs = [p[0] for p in pixs]
        cs = [p[1] for p in pixs]
        img_color[rs, cs] = color

    # 画节点
    for ctr in nodes:
        r_m, c_m = ctr
        cv2.circle(img_color, (int(round(c_m)), int(round(r_m))), 2, (0, 0, 255), -1)

    # 画边
    for i, j in edges:
        r1, c1 = nodes[i]
        r2, c2 = nodes[j]
        p1 = (int(round(c1)), int(round(r1)))
        p2 = (int(round(c2)), int(round(r2)))
        cv2.line(img_color, p1, p2, (255, 0, 0), 1)

    graph = {"nodes": nodes, "edges": edges}
    return img_color, graph


def is_connectivity_preserved(original_mask, resized_img, runnable_color, scale_k):
    """验证缩放后可行域的连通区域数量是否一致"""
    resized_mask = (resized_img == runnable_color).astype(np.uint8)
    original_labels = cv2.connectedComponents(original_mask)[1]
    resized_labels = cv2.connectedComponents(resized_mask)[1]

    # 缩放回原尺寸的标签图
    resized_labels_rescaled = cv2.resize(
        resized_labels,
        (original_mask.shape[1], original_mask.shape[0]),
        interpolation=cv2.INTER_NEAREST,
    )
    # 统计连通区域数量
    original_count = np.max(original_labels)
    resized_count = np.max(resized_labels_rescaled)

    print(f"  >> 连通区域数量：原图 {original_count}，缩放后 {resized_count}")
    return original_count == resized_count


def active_size_map_discretization(img, runnable_color=0, node_radius=30):
    """
    动态调整图像尺寸的地图离散化函数。
    根据可行域最小宽度进行缩放，避免结构断开。
    """
    # 1. 获取可行域的最小宽度
    print("[1/4] 正在计算可行域最小宽度...")

    # 创建二值掩膜：可行域为255，其他为0
    binary_mask = (img == runnable_color).astype(np.uint8) * 255

    # 距离变换：计算每个可行像素到最近障碍物的距离
    dist_transform = cv2.distanceTransform(binary_mask, cv2.DIST_L2, 5)

    # 提取可行域中的距离值
    valid_distances = dist_transform[binary_mask > 0]

    if len(valid_distances) == 0:
        raise ValueError("输入图像中未检测到可行域，请检查 runnable_color 设置。")

    min_distance = np.min(valid_distances)
    median_distance = np.median(valid_distances)
    min_width = min_distance * 2  # 最小宽度为两倍最小距离
    median_width = median_distance * 2  # 中位数宽度

    print(f"  >> 可行域最小宽度: {min_width:.2f} 像素")
    print(f"  >> 可行域中位数宽度: {median_width:.2f} 像素")

    # 2. 根据最小宽度决定缩放比例
    print("[2/4] 正在计算缩放比例...")
    target_width_base = img.shape[1] / 100  # 动态目标宽度
    target_min_width = max(2, int(target_width_base))  # 确保不小于2像素

    scale_width = median_width  # 使用中位数宽度作为缩放依据
    k = target_min_width / scale_width
    print(f"  >> 缩放比例 k = {k:.4f}")

    # 3. 缩放图像并验证连通性
    print("[3/4] 正在缩放图像...")
    new_size = (int(img.shape[1] * k), int(img.shape[0] * k))
    img_shrink = cv2.resize(img, new_size, interpolation=cv2.INTER_AREA)

    # 验证连通性（可选）
    if not is_connectivity_preserved(binary_mask, img_shrink, runnable_color, k):
        print("检测到连通性变化，调整缩放比例...")
        k *= 0.9  # 适当减小缩放比例
        new_size = (int(img.shape[1] * k), int(img.shape[0] * k))
        img_shrink = cv2.resize(img, new_size, interpolation=cv2.INTER_AREA)
        print(f"  >> 缩放比例调整为 k = {k:.4f}")

    # 4. 对缩放后的图像进行离散化处理
    print("[4/4] 正在对缩放后的图像进行离散化处理...")
    img_shrink_color, graph = floating_center_discretize(
        img_shrink, runnable_color, node_radius=int(node_radius * k)
    )

    # 5. 将结果还原至原图尺寸
    print("  >> 还原图像尺寸与图结构...")
    img_r = cv2.resize(
        img_shrink_color, (img.shape[1], img.shape[0]), interpolation=cv2.INTER_CUBIC
    )
    graph["nodes"] = [
        [int(y / k), int(x / k)] for x, y in graph["nodes"]
    ]  # 注意坐标顺序

    return img_r, graph


if __name__ == "__main__":
    img_path = "map.png"  
    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
    img = cv2.resize(img, (600, 400))
    runnable_color = 255
    result_img, graph = active_size_map_discretization(
        img, runnable_color, node_radius=11
    )
    print(
        f"number of node is {len(graph['nodes'])}, number of edge is {len(graph['edges'])}"
    )

    # save
    cv2.imwrite("graph.png", result_img)
    with open("graph.yaml", "w") as f:
        yaml.dump(graph, f)

    # # 保存点坐标
    # with open("points.txt", "w") as f:
    #     for node in graph["nodes"]:
    #         x, y = node
    #         f.write(f"{x} {y}\n")
    # # 生成邻接矩阵并保存
    # num_node = len(graph["nodes"])
    # adj_matrix = np.eye(num_node, dtype=np.int32)
    # for edge in graph["edges"]:
    #     node1_idx = edge[0]
    #     node2_idx = edge[1]
    #     adj_matrix[node1_idx, node2_idx] = True
    #     adj_matrix[node2_idx, node1_idx] = True
    # np.save("adj_matrix.npy", adj_matrix)
    # with open("adj_matrix.txt", "w") as f:
    #     f.write(str(num_node))
    #     f.write("\n")
    #     for i in range(num_node):
    #         for j in range(num_node):
    #             s = 1 if adj_matrix[i, j] else 0
    #             f.write(f"{s} ")
    #         f.write("\n")


    # cv2.imshow("Discretized Map", result_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()